import os
import random
import shutil
import typing
import zipfile
import logging
import datetime
import threading


import interface.utils as Iutils
from IPython.display import display, update_display
from interface.widgets import Widgets, debug
from interface.scripts.train import start as train_start
from interface.scripts.test import start as test_start


logger = logging.getLogger()
Iutils.init_log(logger, ".logs")


class Interface:
    def __init__(self) -> None:
        self.display = display

        self.current_path = "~/.cache"
        self.subfolders = ['torch', 'hub', 'checkpoints']

        self.selected_model = ""

        self.base_path = 'models'
        self.params_filename = "params_config.json"

        self.current_model_dir = ""

        # Param
        self.default_params = {
            "orc_model_flag": "tesseract",
            "confidence_threshold": 0.75,
            "iou_threshold": 0.45
        }

        self.global_params = [
            {"name": key, "value": str(value)}
            for key, value in self.default_params.items()
        ]

        self.params = [*self.global_params]

        self.widgets = Widgets(self)

    def mkdir_subfolders(self):
        for folder in self.subfolders:
            current_path = os.path.join(self.current_path, folder)
            if not os.path.exists(current_path):
                os.makedirs(current_path)

    def find_param(self, name: str):
        for param in self.params:
            if value := param.get(name):
                return value
        if value := self.default_params.get(name):
            return value
        else:
            raise ValueError(
                f"Not param name: {name} in params or default_params")

    def if_dataset_exists(self):
        for folder in ['train', 'val']:
            if os.path.isdir(
                os.path.join(
                    self.base_path,
                    self.selected_model,
                    folder)):
                return True
        return False

    def if_model_exists(self):
        for file in [os.listdir(os.path.join(self.base_path, self.selected_model))]:
            if file.endswith('.pth'):
                return True

    def get_model_folders(self, base_path):
        return sorted(
            [
                name for name in os.listdir(base_path) if os.path.isdir(os.path.join(base_path, name)) and not name.startswith('.')
            ])

    def popup(self, text: str):
        clear_output()
        self.display(Javascript("alert('{}')".format(text)))

    def edit_param(self, index: int):
        self.widget.update_display(
            self.params[index]['name']
        )

    def save_param(self, index: int, value_input):
        if value_input.value:
            if value_input.value.isdigit():
                self.params[index]["value"] = value_input.value
                Iutils.write_params_to_json(
                    os.path.join(self.base_path, self.selected_model,
                                 self.params_filename),
                    [self.params])
                value_input.disabled = True
                return self.widgets.update_display([])
            else:
                text = "Ошибка! Необходимо ввести целое число"
        else:
            text = "Введите значение параметра"
        self.popup(text)

    def on_update_selected_index(self, change):
        if self.selected_model == "":
            if self.widgets.tab.selected_index in (1, 2, 3, 4):
                self.popup("Для начала выберите нейросеть")
                self.widgets.tab.selected_index = 0
        elif not self.if_dataset_exists():
            if self.widgets.tab.selected_index in (2, 3, 4):
                self.popup("Для начала загрузите датасет")
                self.widgets.tab.selected_index = 1
        elif self.widgets.tab.selected_index == 4 and not self.if_model_exists():
            self.popup(
                'Тестирование недоступно. Пожалуйста, обучите нейросеть на вкладке "Обучение"')
            self.widgets.tab.selected_index = 3

    def display_models(self):
        with self.widgets.models_output:
            clear_output()
            self.display(
                self.widgets.generate_models(
                    self.get_model_folders(self.base_path)
                )
            )

    def delete_model(self, model_name):
        model_path = os.path.join(self.base_path, model_name)
        if os.path.exists(model_path):
            shutil.rmtree(model_path)
            if self.selected_model == model_name:
                self.select_model("")
        self.display_models()

    def add_model(self, call):
        logger.info(call)
        new_model_name = self.widgets.new_model_input.value
        if new_model_name:
            new_model_path = os.path.join(
                self.base_path, new_model_name)
            if not os.path.exists(new_model_path):
                os.makedirs(new_model_path)
                Iutils.write_params_to_json(
                    os.path.join(
                        new_model_path, self.params_filename),
                    [*self.global_params]
                )
                logger.info(f"Создана новая модель: {new_model_name}")
            else:
                self.popup(
                    f"Нейросеть с таким именем ({new_model_name}) уже существует!")
            self.widgets.new_model_input.value = ""
        self.display_models()

    def select_model(self, model_name):
        self.selected_model = model_name
        if self.selected_model == "":
            self.widgets.model_label.value = "Нейросеть для обучения: Не выбрана"
            self.widgets.generate_params(self.params)
        self.widgets.fi.file_info = []
        self.widgets.fi.clear()
        self.widgets.imagei.file_info = []
        self.widgets.imagei.clear()
        self.widgets.testing_button.disabled = True
        self.widgets.test_img_preview.value = Iutils.read_file(
            "interface/source/placeholder.jpg", "rb")
        self.widgets.accuracy_label.value = "Качество обучения:  <b>N/A</b>"
        self.widgets.train_time.value = "00:00:00 / 00:00:00 / 0 мин. 0 сек."
        self.widgets.progress_bar.value = 0
        self.widgets.status.style.text_color = "#333333"
        self.widgets.status.value = "Ожидание начала обучения"
        params_path = os.path.join(
            self.base_path, self.selected_model, self.params_filename)
        if os.path.exists(params_path):
            self.params = Iutils.read_params_from_json(params_path)
        else:
            Iutils.write_params_to_json(
                params_path, [*self.global_params]
            )
        self.widgets.update_display([])
        self.widgets.model_label.value = f"Нейросеть для обучения: {self.selected_model}"
        if self.if_dataset_exists():
            self.widgets.uploaded_filenames.value = ""
            self.widgets.uploaded_filenames.value += "Датасет уже загружен<br>"
            for category in [
                d for d in os.listdir(os.path.join(self.base_path, self.selected_model, "train")) if os.path.isdir(
                    os.path.join(self.base_path, self.selected_model, "train", d))]:
                all_train_files = len(os.listdir(os.path.join(
                    self.base_path, self.selected_model, "train", category)))
                all_val_files = len(os.listdir(os.path.join(
                    self.base_path, self.selected_model, "val", category)))
                all_files = all_train_files + all_val_files
                self.widgets.uploaded_filenames.value += f"Категория '{category}': {all_files} файлов<br>"
        else:
            self.widgets.uploaded_filenames.value = ""
            self.widgets.uploaded_filenames.value += "Датасет еще не загружен"
        files = os.listdir(os.path.join(self.base_path, model_name))
        self.widgets.save_model_button.value = '<i class="fa fa-floppy-o" style="margin-left: 40px; margin-right: var(--jp-widgets-inline-margin);"></i><button disabled="true" style="background: none;color: inherit;border: none;padding: 0;font: inherit;cursor: pointer;outline: inherit;">Модель еще не обучена</button>'
        if files := os.listdir(os.path.join(self.base_path, model_name)):
            for file in files:
                if file.endswith('.pth'):
                    self.create_model_link(True)
                    break
        self.widgets.generate_params(self.params)

    def create_model_link(self, on_select_loading):
        model_dir = os.path.abspath(
            os.path.join(
                self.base_path,
                self.selected_model))
        folder_path = os.path.join(
            self.base_path, self.selected_model)
        max_date, max_file = None, None
        for file_name in os.listdir(model_dir):
            if file_name.startswith(self.selected_model) and file_name.endswith(".pth"):
                file_path = os.path.join(model_dir, file_name)
                file_stat = os.stat(file_path)
                if max_date is None or file_stat.st_mtime > max_date:
                    max_date = file_stat.st_mtime
                    max_file = file_name
        if max_file:
            config = Iutils.load_config(
                os.path.join(folder_path, "train_config.yaml")
            )
            max_file_size = max_file.split(
                ".")[0] + f"_size_{int(config['image_size'])}" + "." + max_file.split(".")[1]
            os.rename(os.path.join(model_dir, max_file),
                      os.path.join(model_dir, max_file_size))
            required_files = ['classes_config.yaml', max_file_size]

            missing_files = [f for f in required_files if not os.path.isfile(
                os.path.join(folder_path, f))]

            if not missing_files:
                archive_path = os.path.join(
                    model_dir, f"{self.selected_model}.zip")
                files_to_archive = [os.path.join(
                    folder_path, f) for f in required_files]
                if not on_select_loading:
                    if os.path.exists(archive_path):
                        os.remove(archive_path)
                    Iutils.write_zip_file(files_to_archive, archive_path)
                os.rename(os.path.join(model_dir, max_file_size),
                          os.path.join(model_dir, max_file))
                username = os.environ['JUPYTERHUB_USER']
                link = f'{os.environ["JUPYTERHUB_BASE_URL"]}/user/{username}/files/{os.path.join(*archive_path.split("/")[4:])}'
                self.widgets.save_model_button.value = f"""
                    <i class="fa fa-floppy-o" style="margin-left: 32px; margin-right: var(--jp-widgets-inline-margin);"></i>
                    <button onclick="window.location.href = window.location.origin + \'{link}\'" style="background: none;color: inherit;border: none;padding: 0;font: inherit;cursor: pointer;outline: inherit;">Сохранить обученную модель</button>
                """
            else:
                print("Отсутствуют следующие файлы:", missing_files)

    def on_test_upload(self, change):
        self.widgets.category_label.value = "Полученный тип документа: "
        if isinstance(self.widgets.imagei.file_info, list):
            self.widgets.testing_button.disabled = True
            self.widgets.test_img_preview.value = open(
                "interface/source/placeholder.jpg", "rb").read()
            return
        file = self.widgets.imagei.get_files()[0]
        model_dir = os.path.abspath(
            os.path.join(self.base_path, self.selected_model))
        test_photo_path = os.path.join(model_dir, "fortesting_img.jpg")
        with open(test_photo_path, 'wb') as f:
            f.write(file['file_obj'].read())
        self.widgets.test_img_preview.value = open(
            test_photo_path, "rb").read()
        self.widgets.testing_button.disabled = False

    @debug.capture(True)
    def on_upload_change(self, change):
        if not self.widgets.fi.file_info:
            logger.info(f"fi.file_info: {self.widgets.fi.file_info}")
            return
        self.widgets.uploaded_filenames.value = ""
        if self.if_dataset_exists():
            self.widgets.uploaded_filenames.value = ""
            shutil.rmtree(
                os.path.join(self.base_path, self.selected_model))
            # self.popup("Внимание! Датасет будет перезаписан")
        file = self.widgets.fi.get_files()[0]
        data = Iutils.upload_dataset(
            os.path.join(
                self.base_path, self.selected_model),
            file
        )

        for _type, category, count in data:
            self.widgets.uploaded_filenames.value += f"Тип: '{_type}' Категория: '{category}' загружена: {len(count)} файлов<br>"

    # Train

    def create_train_config_yaml(self, model_dir):
        Iutils.create_config_model(
            os.path.join(model_dir, 'train_config.yaml'),
            {
                'data_dir': model_dir,
                'output_classes_path': os.path.join(model_dir, "classes_config.yaml"),
                'image_size': self.find_param("image_size"),
                'batch_size': self.find_param("batch_size"),
                'num_epochs': self.find_param("num_epochs"),
                'model_save_path': os.path.join(
                    model_dir, self.selected_model),
            }
        )

    def train_model(self, start_time: datetime.datetime):
        # Widgets
        self.widgets.status.value = f"Осталось 0 часов 0 минут 0 секунд / Выполнено  ({self.widgets.progress_bar.value * 100:.0f}%)"
        self.widgets.accuracy_label.value = "Качество обучения:  <b>N/A</b>"
        # path_to_return = os.getcwd()
        # os.chdir(current_model_dir)
        try:
            train_start(
                self.widgets.progress_bar,
                self.widgets.status,
                self.widgets.accuracy_label
            )
        except Exception as error:
            logger.error(error)

        # os.chdir(path_to_return)
        # Widgets
        self.widgets.status.style.text_color = "#333333"
        if self.train_stop:
            # Widgets
            self.widgets.status.value = "Статус: Обучение остановлено"
            self.widgets.status.style.text_color = "#E22C38"
        else:
            self.widgets.status.value = "Статус: Обучение завершено"
            self.widgets.status.style.text_color = "green"

        self.widgets.train_time.value = Iutils.train_time(start_time)

        # todo вынести в отдельный метод, тк используется в create yaml
        # create_model_link(on_select_loading=False)
        self.widgets.train_button.disabled = False
        self.widgets.stop_button.disabled = True

    def start_training(self, call):
        self.current_model_dir = os.path.abspath(
            os.path.join(self.base_path, self.selected_model))
        self.create_train_config_yaml(self.current_model_dir)
        start_time = datetime.datetime.now()
        self.train_stop = False
        # Widgets
        self.widgets.progress_bar.value = 0
        self.widgets.train_time.value = f"{start_time.strftime('%H:%M:%S')} / 00:00:00 / 0 мин. 0 сек."
        self.widgets.train_button.disabled = True
        self.widgets.stop_button.disabled = True
        self.widgets.status.style.text_color = "#333333"

        training_thread = threading.Thread(target=self.train_model(start_time))
        training_thread.start()

    def stop_training(self, call):
        self.train_stop = True
        self.widgets.train_button.disabled = False
        self.widgets.stop_button.disabled = True

    # Test
    def create_test_config_yaml(self, model_dir):
        train_config = Iutils.load_config(
            os.path.join(model_dir, 'train_config.yaml'))

        model_save_path = ""
        max_date, max_file = None, None
        for file_name in os.listdir(model_dir):
            if file_name.startswith(self.selected_model) and file_name.endswith(".pth"):
                file_stat = os.stat(os.path.join(model_dir, file_name))
                if max_date is None or file_stat.st_mtime > max_date:
                    max_date = file_stat.st_mtime
                    max_file = file_name
        if max_file:
            model_save_path = os.path.join(model_dir, max_file)
        else:
            print("Err! Модель не найдена")

        Iutils.create_config_model(
            os.path.join(model_dir, 'test_config.yaml'),
            {
                'image_size': train_config.get('image_size'),
                'model_weights_path': model_save_path,
                'image_path': os.path.join(model_dir, 'fortesting_img.jpg'),
                'class_config_path': train_config.get('output_classes_path')
            }

        )

    def start_testing(self, call):
        # current_model_dir = os.path.abspath(
        #     os.path.join(base_path, selected_model))
        self.create_test_config_yaml(self.current_model_dir)
        self.test_model()

    def test_model(self):
        # path_to_return = os.getcwd()
        # os.chdir(current_model_dir)
        try:
            predicted_class = test_start()
            self.widgets.category_label.value = f"Полученный тип документа: {predicted_class}"
        except Exception as e:
            print(e)
        # os.chdir(path_to_return)

    def __call__(self, *args: typing.Any, **kwds: typing.Any) -> typing.Any:
        # Style
        self.display(
            self.widgets.build_style(
                raw_style=Iutils.read_file("interface/source/style.css", "r")
            )
        )
       
        # Call
        self.display(self.widgets.tab)
        self.widgets.tab_observe()
        self.display_models()
        self.display(self.widgets.nav_buttons)
        self.display(self.widgets.notify_output)
